
#include "cuda_runtime.h"
#include "device_launch_parameters.h"

#include "measure.hpp"

#include <stdio.h>
#include <iostream>
#include <string>
#include <conio.h>
#include <cstdlib>

#include <cublas_v2.h>

using namespace std;

typedef float mat_field_t;

struct MatMulData
{
    int m, n, z;
    mat_field_t *h_c;
    mat_field_t *h_a;
    mat_field_t *h_b;
    mat_field_t *d_c;
    mat_field_t *d_a;
    mat_field_t *d_b;
};

cudaError_t chooseSmartestDevice();

cublasHandle_t handle;

cublasStatus_t ret;

const float alpha = 1.0f;
const float beta  = 0.0f;

void cleanupMatrices(int size, void** pData)
{
    MatMulData *pMatrices = (MatMulData*)*pData;
    bool worked = true;
    for(int test=0; test<5 ; ++test)
    {
        //check for correctness
        int i=rand()%size;
        int j=rand()%size;
        mat_field_t result = 0;
        for(int o=0; o<size ; ++o)
        {
            //cout << ": " << pMatrices->h_a[i+o*size] << "*" << pMatrices->h_a[o+j*size] << endl;
            result += pMatrices->h_b[i+o*size]*pMatrices->h_a[o+j*size];
        }
        const mat_field_t expected = pMatrices->h_c[i+j*size];
        const mat_field_t tolerance = 0.0001;
        if(result>(expected+tolerance)||result<(expected-tolerance))
        {
            cout << "Did not calculate correctly!" << result << " != " << expected << "(i: " << i << ", " << j << ")" << endl;
            worked = false;
        }
    }
    if(!worked)
    {
        cout << "did not work!" << endl;
    }
    delete [] pMatrices->h_c;
    delete [] pMatrices->h_a;
    delete [] pMatrices->h_b;
    delete pMatrices;
}

void mulMatrices(int size, void** pData)
{
    MatMulData *pMatrices = (MatMulData*)*pData;
    
    const int m = pMatrices->m;
    const int n = pMatrices->n;
    const int z = pMatrices->z;
    const size_t fieldSize = sizeof(mat_field_t);
    
    cudaError_t cudaStatus = cudaMalloc((void**)&pMatrices->d_c, m*n*fieldSize);
    if (cudaStatus != cudaSuccess) {
        size_t avail;
        size_t total;
        cudaMemGetInfo(&avail, &total);
        size_t tried = m*n*fieldSize;
        fprintf(stderr, "cudaMalloc failed! avail: %u total %u tried: \n", avail, total, tried);
        cudaDeviceReset();
        throw "error";
    }

    cudaStatus = cudaMalloc((void**)&pMatrices->d_a, m*z*fieldSize);
    if (cudaStatus != cudaSuccess) {
        size_t avail;
        size_t total;
        cudaMemGetInfo(&avail, &total);
        size_t tried = m*z*fieldSize;
        fprintf(stderr, "cudaMalloc failed! avail: %u total %u tried: %u (%s)\n", avail, total, tried, cudaGetErrorString(cudaGetLastError()));
        cudaDeviceReset();
        throw "error";
    }

    cudaStatus = cudaMalloc((void**)&pMatrices->d_b, n*z*fieldSize);
    if (cudaStatus != cudaSuccess) {
        size_t avail;
        size_t total;
        cudaMemGetInfo(&avail, &total);
        size_t tried = n*z*fieldSize;
        fprintf(stderr, "cudaMalloc failed! avail: %u total %u tried: %u (%s)\n", avail, total, tried, cudaGetErrorString(cudaGetLastError()));
        throw "error";
    }


    cudaStatus = cudaMemcpy(pMatrices->d_a, pMatrices->h_a, m*z*fieldSize, cudaMemcpyHostToDevice);
    if (cudaStatus != cudaSuccess) {
        fprintf(stderr, "cudaMemcpy failed! (%s)\n", cudaGetErrorString(cudaGetLastError()));
        cudaDeviceReset();
        throw "error";
    }

    cudaStatus = cudaMemcpy(pMatrices->d_b, pMatrices->h_b, n*z*fieldSize, cudaMemcpyHostToDevice);
    if (cudaStatus != cudaSuccess) {
        fprintf(stderr, "cudaMemcpy failed! (%s)\n", cudaGetErrorString(cudaGetLastError()));
        cudaDeviceReset();
        throw "error";
    }

    ret = cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, z, &alpha, pMatrices->d_b, size, pMatrices->d_a, size, &beta, pMatrices->d_c, size);

    if (ret != CUBLAS_STATUS_SUCCESS)
    {
        printf("cublasSgemm returned error code %d, line(%d)\n", ret, __LINE__);
        exit(EXIT_FAILURE);
    }

    cudaStatus = cudaGetLastError();
    if (cudaStatus != cudaSuccess) {
        fprintf(stderr, "after kernel execution! (%s)\n", cudaGetErrorString(cudaGetLastError()));
        cudaDeviceReset();
        throw "error";
    }
    cudaStatus = cudaThreadSynchronize();
    if (cudaStatus != cudaSuccess) {
        fprintf(stderr, "sync after kernel execution! (%s)\n", cudaGetErrorString(cudaGetLastError()));
        cudaDeviceReset();
        throw "error";
    }
    cudaStatus = cudaGetLastError();
    if (cudaStatus != cudaSuccess) {
        fprintf(stderr, "after sync kernel execution! (%s)\n", cudaGetErrorString(cudaGetLastError()));
        cudaDeviceReset();
        throw "error";
    }
    cudaStatus = cudaMemcpy(pMatrices->h_c, pMatrices->d_c, pMatrices->m*pMatrices->n*sizeof(mat_field_t), cudaMemcpyDeviceToHost);
    if (cudaStatus != cudaSuccess) {
        fprintf(stderr, "cudaMemcpy BACK failed! (%s)\n", cudaGetErrorString(cudaGetLastError()));
        cudaDeviceReset();
        throw "error";
    }
    cudaFree(pMatrices->d_c);
    cudaFree(pMatrices->d_a);
    cudaFree(pMatrices->d_b);
    cudaStatus = cudaThreadSynchronize();
    if (cudaStatus != cudaSuccess) {
        fprintf(stderr, "cudaThreadSynchronize failed!\n");
        cudaDeviceReset();
        throw "error";
    }
}

__host__ void initMatrices(int size, void** pData)
{
    MatMulData* pMatrices = new MatMulData();
    pMatrices->m = size;
    pMatrices->n = size;
    pMatrices->z = size;
    const int m = pMatrices->m;
    const int n = pMatrices->n;
    const int z = pMatrices->z;

    pMatrices->h_a = new mat_field_t[m*z];
    for(int i=0 ; i<m*z ; ++i)
    {
        pMatrices->h_a[i] = i*0.1;
    }
    pMatrices->h_b = new mat_field_t[n*z];
    for(int i=0 ; i<n*z ; ++i)
    {
        pMatrices->h_b[i] = 1;
    }

    //cout << "Allocating 3 * " << m*n << " * " << fieldSize << " = " << (m*z+n*z+m*n)*fieldSize << " byte. 2GB=" << 2ul*1024ul*1204ul*1024ul << endl;

    const int matSize = pMatrices->m*pMatrices->n;
    pMatrices->h_c = new mat_field_t[matSize];

    *pData = pMatrices;
}

int main()
{
    chooseSmartestDevice();
    Measure measure;
    measure.setOutstream("zeitmessungCublas.txt");
    measure.setDebugOut(cout);
    measure.addFunction("cublasMatMul", mulMatrices, initMatrices, cleanupMatrices);
    measure.addNewMeasureInterval(200, 11, 3000);
    measure.setRepetitions(3);

    //// Init cublas
    ret = cublasCreate(&handle);

    if (ret != CUBLAS_STATUS_SUCCESS)
    {
        fprintf(stderr, "cublasCreate returned error code %d, line(%d)\n", ret, __LINE__);
        throw "error";
    }

    measure.run();
    
    ret = cublasDestroy(handle);

    if (ret != CUBLAS_STATUS_SUCCESS)
    {
        fprintf(stderr, "cublasDestroy returned error code %d, line(%d)\n", ret, __LINE__);
        throw "error";
    }

    getch();
}

cudaError_t chooseSmartestDevice()
{
    cudaError_t cudaStatus;

    int deviceId;
    string deviceName;
    int deviceCount = 0;
    cudaError_t error_id = cudaGetDeviceCount(&deviceCount);

    if (error_id != cudaSuccess)
    {
        printf("cudaGetDeviceCount returned %d\n-> %s\n", (int)error_id, cudaGetErrorString(error_id));
        exit(EXIT_FAILURE);
    }

    // This function call returns 0 if there are no CUDA capable devices.
    if (deviceCount == 0)
    {
        printf("There are no available device(s) that support CUDA\n");
    }
    else
    {
        printf("Detected %d CUDA Capable device(s)\n", deviceCount);
    }

    int major = 0;
    int minor = 0;
    for (int dev = 0; dev < deviceCount; ++dev)
    {
        cudaDeviceProp deviceProp;
        cudaGetDeviceProperties(&deviceProp, dev);

        printf("\nDevice %d: \"%s\"\n", dev, deviceProp.name);
        if(deviceProp.major > major || (deviceProp.major == major && deviceProp.minor > minor))
        {
            major = deviceProp.major;
            minor = deviceProp.minor;
            deviceId = dev;
            deviceName = deviceProp.name;
        }
        //if(strcmpi(deviceProp.name, "GeForce GTX 260") == 0)
        //{
        //	major = deviceProp.major;
        //	minor = deviceProp.minor;
        //	deviceId = dev;
        //	deviceName = deviceProp.name;
        //	break;
        //}
    }
    printf("\nCoosing your coolest CUDA device: %d: \"%s\" (highest compute capability: %d.%d)\n", deviceId, deviceName.c_str(), major, minor);
    // Choose which GPU to run on, change this on a multi-GPU system.
    cudaStatus = cudaSetDevice(deviceId);
    if (cudaStatus != cudaSuccess) {
        fprintf(stderr, "cudaSetDevice failed!  Do you have a CUDA-capable GPU installed?");
    }
    return cudaStatus;
}